from vtkplotter import *
import vtk
import argparse
import numpy as np
import os, sys, math
from vtk.util import numpy_support
from scipy.stats import f

def normalize(vec):
    d = np.linalg.norm(vec)
    if d > 0:
        return vec / d
    return vec

def kalman(current, last):
    if last > current:
        return last
    else:
        return (last * 7 + current) / 8

def fit_line(points, normal, margin=0, last_width=None):
    center = points.mean(axis=0)
    uu, dd, vv = np.linalg.svd(points - center)
    vv = vv[0] / np.linalg.norm(vv[0])
    vv = normalize(vv - normal * np.dot(vv, normal))
    # vv contains the first principal component, i.e. the direction
    # vector of the best fit line in the least squares sense.
    l = np.dot(points - center, vv)
    half_width = (l.max() - l.min()) / 2 + margin
    if last_width:
        half_width = kalman(half_width, last_width)
    
    center = center + vv * (l.max() + l.min()) / 2
    p1 = center - vv * half_width
    p2 = center + vv * half_width
    return p1, p2, center, half_width

def fit_line2(points, direction, margin=0, last_width=None):
    center = points.mean(axis=0)
    l = np.dot(points - center, direction)
    half_width = (l.max() - l.min()) / 2 + margin
    if last_width:
        half_width = kalman(half_width, last_width)
    
    center = center + direction * (l.max() + l.min()) / 2
    p1 = center - direction * half_width
    p2 = center + direction * half_width
    return p1, p2, center, half_width

def fit_plane(points, scale):
    datamean = points.mean(axis=0)
    res = np.linalg.svd(points - datamean)
    dd, vv = res[1], res[2]
    xyz_min = points.min(axis=0)
    xyz_max = points.max(axis=0)
    s = np.linalg.norm(xyz_max - xyz_min) * scale
    n = np.cross(vv[0], vv[1])
    pla = Plane(datamean, n, s, s)
    pla.info["normal"] = n
    pla.info["center"] = datamean
    pla.info["variance"] = dd[2]
    return pla

def actor_to_volume(actor, bounds, spacing=(1, 1, 1)):
    pd = actor.polydata()

    whiteImage = vtk.vtkImageData()
    # bounds = pd.GetBounds()

    whiteImage.SetSpacing(spacing)

    # compute dimensions
    dim = [0, 0, 0]
    for i in [0, 1, 2]:
        dim[i] = int(np.ceil((bounds[i * 2 + 1] - bounds[i * 2]) / spacing[i]))
    whiteImage.SetDimensions(dim)
    whiteImage.SetExtent(0, dim[0] - 1, 0, dim[1] - 1, 0, dim[2] - 1)

    origin = bounds[0::2]
    whiteImage.SetOrigin(origin)
    whiteImage.AllocateScalars(vtk.VTK_UNSIGNED_CHAR, 1)

    # fill the image with foreground voxels:
    inval = 255
    count = whiteImage.GetNumberOfPoints()
    whiteArray = np.zeros(count, dtype=np.uint8)
    whiteArray[:] = inval
    whiteImage.GetPointData().SetScalars(numpy_support.numpy_to_vtk(whiteArray))

    # polygonal data --> image stencil:
    pol2stenc = vtk.vtkPolyDataToImageStencil()
    pol2stenc.SetInputData(pd)
    pol2stenc.SetOutputOrigin(origin)
    pol2stenc.SetOutputSpacing(spacing)
    pol2stenc.SetOutputWholeExtent(whiteImage.GetExtent())
    pol2stenc.Update()

    # cut the corresponding white image and set the background:
    outval = 0
    imgstenc = vtk.vtkImageStencil()
    imgstenc.SetInputData(whiteImage)
    imgstenc.SetStencilConnection(pol2stenc.GetOutputPort())
    imgstenc.ReverseStencilOff()
    imgstenc.SetBackgroundValue(outval)
    imgstenc.Update()
    return Volume(imgstenc.GetOutput())

def pcaPoints(_points):    
    cov = np.cov(_points, rowvar=0)     # covariance matrix
    U, s, R = np.linalg.svd(cov)  # singular value decomposition
    p, n = s.size, _points.shape[0]
    fppf = f.ppf(0.95, p, n-p)*(n-1)*p*(n+1)/n/(n-p)  # f % point function
    ua, ub, uc = np.sqrt(s*fppf)*2  # semi-axes (largest first)
    _center = np.mean(_points, axis=0)     # centroid of the hyperellipsoid

    _main_axis = np.array((R[0][0] * ua, R[0][1] * ua, R[0][2] * ua))
    _bi_axis = np.array((R[1][0] * ub, R[1][1] * ub, R[1][2] * ub))
    _normal = np.array((R[2][0] * uc, R[2][1] * uc, R[2][2] * uc))

    return _center, _main_axis, _bi_axis, _normal

def compute_u_plate(upper_poly, lower_poly, sample_num=20, margin=3, plane_offset=0, resolution=0.2):
    upper_actor = Actor(upper_poly)
    lower_actor = Actor(lower_poly)
    upper_deci = upper_actor.decimate(5000 / upper_actor.N()).clean()
    upper_points = upper_deci.getPoints()
    lower_deci = lower_actor.decimate(5000 / lower_actor.N()).clean()
    lower_points = lower_deci.getPoints()

    upper_plane = fit_plane(upper_points, 2).alpha(0.5)
    lower_plane = fit_plane(lower_points, 2).alpha(0.5)
    center = (upper_plane.info['center'] + lower_plane.info['center']) / 2

    if np.dot(upper_plane.info['normal'], center - upper_plane.info['center']) < 0:
        upper_plane.info['normal'] = -upper_plane.info['normal']

    upper_points = upper_points[np.dot(upper_points - upper_plane.info['center'], upper_plane.info['normal']) >= 0]

    if np.dot(lower_plane.info['normal'], center - lower_plane.info['center']) < 0:
        lower_plane.info['normal'] = -lower_plane.info['normal']

    lower_points = lower_points[np.dot(lower_points - lower_plane.info['center'], lower_plane.info['normal']) >= 0]
    
    points = np.concatenate((upper_points, lower_points), axis=0)
    
    cov = np.cov(points, rowvar=0)     # covariance matrix
    U, s, R = np.linalg.svd(cov)  # singular value decomposition
    p, n = s.size, points.shape[0]
    fppf = f.ppf(0.95, p, n-p)*(n-1)*p*(n+1)/n/(n-p)  # f % point function
    ua, ub, uc = np.sqrt(s*fppf)*2  # semi-axes (largest first)
    center = np.mean(points, axis=0)     # centroid of the hyperellipsoid

    main_axis = np.array((R[0][0] * ua, R[0][1] * ua, R[0][2] * ua))
    bi_axis = np.array((R[1][0] * ub, R[1][1] * ub, R[1][2] * ub))
    normal = np.array((R[2][0] * uc, R[2][1] * uc, R[2][2] * uc))
    
    points = points - center
    main_axis = normalize(main_axis)
    bi_axis = normalize(bi_axis)
    x = np.dot(points, main_axis)
    y = np.dot(points, bi_axis)
    z = np.dot(points, normal)
    theta = np.arctan2(y, x)
    hist, bin_edges = np.histogram(theta, bins=sample_num, range=(-math.pi, math.pi))
    hist_mean = hist.mean()
    hist_std = hist.std()

    samples = []
    main_axis = main_axis.reshape(1, 3)
    bi_axis = bi_axis.reshape(1, 3)
    start = 0
    for i in range(hist.shape[0]):
        if hist[i] <= hist_mean - hist_std:
            if hist[(i + 1) % hist.shape[0]] > hist_mean - hist_std:
                start = len(samples)
            continue

        sample = (theta >= bin_edges[i]) * (theta < bin_edges[i + 1])
        num = np.sum(sample)
        sample = x[sample].reshape((num, 1)) * main_axis + y[sample].reshape((num, 1)) * bi_axis
        sample = sample.mean(axis=0)
        samples.append(sample)

    samples = np.array(samples)
    sorted_samples = []
    for i in range(samples.shape[0]):
        sorted_samples.append(samples[(start + i) % samples.shape[0]])
    
    middle_line = sorted_samples[len(sorted_samples) // 2]
    middle_line = normalize(middle_line)
    main_axis = main_axis.reshape(3)
    bi_axis = bi_axis.reshape(3)
    if abs(np.dot(middle_line, main_axis)) > abs(np.dot(middle_line, bi_axis)):
        if np.dot(middle_line, main_axis) < 0:
            main_axis = -main_axis
            x = -x
            theta = np.arctan2(y, x)
        middle_line = main_axis * ua / 2
    else:
        
        if np.dot(middle_line, bi_axis) < 0:
            bi_axis = -bi_axis
            y = -y
        
        theta = np.arctan2(x, y)
        middle_line = bi_axis * ub / 2

    main_arrow = Line(center, center + middle_line, lw=4).color('purple')

    # sample = (theta >= -math.pi / 2) * (theta < math.pi / 2)
    # num = np.sum(sample)
    # arc = x[sample].reshape((num, 1)) * main_axis + y[sample].reshape((num, 1)) * bi_axis
    arc_samples = []
    arc_outter = []
    arc_inner = []
    half_width = None    
    for i in range(sample_num // 2):
        angle0 = math.pi * (-0.5 + 2 * i / sample_num)
        angle1 = math.pi * (-0.5 + 2 * (i + 1) / sample_num)
        sample = (theta >= angle0) * (theta < angle1)
        # min_z = z[sample].min()
        # max_z = z[sample].max()
        # mean_z = (max_z + min_z) / 2
        # quad_z_size = (max_z - min_z) / 4
        # sample = sample * (z >= (mean_z - quad_z_size)) * (z <= (mean_z + quad_z_size)) 
        num = np.sum(sample)
        sample = x[sample].reshape((num, 1)) * main_axis + y[sample].reshape((num, 1)) * bi_axis
        p0, p1, line_center, half_width = fit_line(sample, normal, margin=margin, last_width=half_width)
        arc_samples.append(line_center)
        if np.linalg.norm(p0) < np.linalg.norm(p1):
            arc_inner.append(p0)
            arc_outter.append(p1)
        else:
            arc_inner.append(p1)
            arc_outter.append(p0)

    left_side_sample = theta < -math.pi / 2    
    num = np.sum(left_side_sample)
    # left_side = x[left_side_sample].reshape((num, 1)) * main_axis + y[left_side_sample].reshape((num, 1)) * bi_axis
    # _, left_direction, _, _ = pcaPoints(left_side)
    # left_direction = normalize(left_direction)

    right_side_sample = theta> math.pi / 2
    num = np.sum(right_side_sample)
    # right_side = x[right_side_sample].reshape((num, 1)) * main_axis + y[right_side_sample].reshape((num, 1)) * bi_axis
    # _, right_direction, _, _ = pcaPoints(right_side)
    # right_direction = normalize(right_direction)

    # left_binorm = np.cross(left_direction, normal)
    # right_binorm = np.cross(right_direction, normal)
    middle_axis = normalize(middle_line)
    new_y = np.dot(points, middle_axis)
    left_y = new_y[left_side_sample]
    right_y = new_y[right_side_sample]
    min_y = min(left_y.min(), right_y.min())
    max_y = max(left_y.max(), right_y.max())
    step_y = (max_y - min_y) / (sample_num // 2)
    left_samples = []
    left_outter = []
    left_inner = []
    right_samples = []
    right_outter =[]
    right_inner = []
    left_half_width = half_width
    right_half_width = None
    direction = normalize(np.cross(middle_axis, normal))
    for i in range(3, sample_num // 2):
        y0 = min_y + i * step_y
        y1 = min_y + (i + 1) * step_y
        sample = left_side_sample * (new_y >= y0) * (new_y < y1)
               
        num = np.sum(sample) 
        sample = x[sample].reshape((num, 1)) * main_axis + y[sample].reshape((num, 1)) * bi_axis
        p0, p1, line_center, left_half_width = fit_line2(sample, direction, margin=margin, last_width=left_half_width)
        left_samples.append(line_center)
        if np.linalg.norm(p0) < np.linalg.norm(p1):
            left_inner.append(p0)
            left_outter.append(p1)
        else:
            left_inner.append(p1)
            left_outter.append(p0)

        sample = right_side_sample * (new_y >= y0) * (new_y < y1)

        num = np.sum(sample)
        sample = x[sample].reshape((num, 1)) * main_axis + y[sample].reshape((num, 1)) * bi_axis
        p0, p1, line_center, right_half_width = fit_line2(sample, direction, margin=margin, last_width=right_half_width)
        right_samples.append(line_center)
        if np.linalg.norm(p0) < np.linalg.norm(p1):
            right_inner.append(p0)
            right_outter.append(p1)
        else:
            right_inner.append(p1)
            right_outter.append(p0)

    
    right_samples.reverse()
    sorted_samples = left_samples
    sorted_samples.extend(arc_samples)
    sorted_samples.extend(right_samples)
    sorted_samples = np.array(sorted_samples) + center
    
    u_silhouette = left_outter[2:]
    u_silhouette.extend(arc_outter)
    right_outter.reverse()
    u_silhouette.extend(right_outter[:-2])
    u_silhouette.extend(right_inner[2:])
    arc_inner.reverse()
    u_silhouette.extend(arc_inner)
    left_inner.reverse()
    u_silhouette.extend(left_inner[:-2])
    
    # u_silhouette.append((u_silhouette[0] + u_silhouette[-1]) / 2)
    u_silhouette.extend(u_silhouette[0:4])
    u_silhouette = np.array(u_silhouette) + center
    u_silhouette = Spline(u_silhouette, res=u_silhouette.shape[0]).getPoints()[2:-2, :]
    u_silhouette_spline = []
    for i in range(3):
        u_silhouette_spline.append(vtk.vtkCardinalSpline())
        u_silhouette_spline[i].ClosedOn()
        u_silhouette_spline[i].SetParametricRange(0, u_silhouette.shape[0])
    for i in range(u_silhouette.shape[0]):
        for j in range(3):
            u_silhouette_spline[j].AddPoint(i, u_silhouette[i, j])

    u_sample_num = 1000
    u_silhouette_poly = vtk.vtkPolyData()
    ppoints = vtk.vtkPoints()
    lines = vtk.vtkCellArray()  # Create the polyline
    lines.InsertNextCell(u_sample_num + 1)
    for i in range(u_sample_num):
        pt = np.zeros(3)
        for j in range(3):
            pt[j] = u_silhouette_spline[j].Evaluate(i * u_silhouette.shape[0]/ u_sample_num)
        ppoints.InsertNextPoint(pt)
        lines.InsertCellPoint(i)
    lines.InsertCellPoint(0)
    u_silhouette_poly.SetPoints(ppoints)
    # u_silhouette_poly.SetLines(lines)
    u_silhouette_poly.SetPolys(lines)
    triangle = vtk.vtkTriangleFilter()
    triangle.SetPassLines(True)
    triangle.SetPassVerts(True)
    triangle.SetInputData(u_silhouette_poly)
    triangle.Update()
    
    transformation = vtk.vtkTransform()
    transformation.Translate(-normal)
    tf = vtk.vtkTransformPolyDataFilter()
    tf.SetTransform(transformation)
    tf.SetInputConnection(triangle.GetOutputPort())
    tf.Update()
    extrude = vtk.vtkLinearExtrusionFilter()
    extrude.SetExtrusionTypeToVectorExtrusion()
    extrude.SetVector(normal * 2)
    extrude.SetInputConnection(tf.GetOutputPort())
    extrude.Update()
  
    u_actor = Actor(extrude.GetOutput())
    u_actor.computeNormals()

    transformation = vtk.vtkTransform()
    transformation.Translate(-upper_plane.info['normal'] * plane_offset)
    upper_plane.transformMesh(transformation)
    extrude = vtk.vtkLinearExtrusionFilter()
    extrude.SetExtrusionTypeToVectorExtrusion()
    extrude.SetVector(-upper_plane.info['normal'])
    extrude.SetScaleFactor(np.linalg.norm(normal) * 2)
    extrude.SetInputData(upper_plane.polydata())
    extrude.Update()
    upper_box = Actor(extrude.GetOutput()).alpha(0.5).color('red')
    upper_box.computeNormals()

    
    transformation = vtk.vtkTransform()
    transformation.Translate(-lower_plane.info['normal'] * plane_offset)
    lower_plane.transformMesh(transformation)
    extrude = vtk.vtkLinearExtrusionFilter()
    extrude.SetExtrusionTypeToVectorExtrusion()
    extrude.SetVector(-lower_plane.info['normal'])
    extrude.SetScaleFactor(np.linalg.norm(normal) * 2)
    extrude.SetInputData(lower_plane.polydata())
    extrude.Update()
    lower_box = Actor(extrude.GetOutput()).alpha(0.5).color('blue')
    lower_box.computeNormals()

    u_bounds = np.array(u_actor.bounds())
    upper_bounds = np.array(upper_actor.bounds())
    lower_bounds = np.array(lower_actor.bounds())
    spacing = [resolution, resolution, resolution]
    bounds=np.zeros(6)
    bounds[0] = min(u_bounds[0], min(upper_bounds[0], lower_bounds[0])) - spacing[0]
    bounds[2] = min(u_bounds[2], min(upper_bounds[2], lower_bounds[2])) - spacing[1]
    bounds[4] = min(u_bounds[4], min(upper_bounds[4], lower_bounds[4])) - spacing[2]
    bounds[1] = max(u_bounds[1], max(upper_bounds[1], lower_bounds[1])) + spacing[0]
    bounds[3] = max(u_bounds[3], max(upper_bounds[3], lower_bounds[3])) + spacing[1]
    bounds[5] = max(u_bounds[5], max(upper_bounds[5], lower_bounds[5])) + spacing[2]
    u_volume = actor_to_volume(u_actor, bounds=bounds, spacing=spacing)
    upper_box_volume = actor_to_volume(upper_box, bounds=bounds, spacing=spacing)
    u_volume = volumeOperation(u_volume, '-', upper_box_volume)
    lower_box_volume = actor_to_volume(lower_box, bounds=bounds, spacing=spacing)
    u_volume = volumeOperation(u_volume, '-', lower_box_volume)
    upper_volume = actor_to_volume(upper_actor, bounds=bounds, spacing=spacing)
    u_volume = volumeOperation(u_volume, '-', upper_volume)
    lower_volume = actor_to_volume(lower_actor, bounds=bounds, spacing=spacing)
    u_volume = volumeOperation(u_volume, '-', lower_volume)
    u_actor = u_volume.isosurface().smoothLaplacian(relaxfact=0.5, edgeAngle=45, featureAngle=90).fillHoles()
    return u_actor.polydata()
    
def onKeyPress(key):
    global vp, upper_actor, lower_actor, u_actor, margin, plane_offset
    key = key.lower()
    if key =='a':
        if u_actor:
            if u_actor.alpha() == 1.0:
                u_actor.alpha(0.5)
            else:
                u_actor.alpha(1)
        else:
            if upper_actor.alpha() == 1.0:
                upper_actor.alpha(0.3)
                lower_actor.alpha(0.3)
            else:
                upper_actor.alpha(1)
                lower_actor.alpha(1)
        vp.interactor.Render()
        return

    if key =='v':
        if upper_actor.alpha() > 0.0:
            upper_actor.alpha(0.0)
            lower_actor.alpha(0.0)
        else:
            upper_actor.alpha(1.0)
            lower_actor.alpha(1.0)
        vp.interactor.Render()
        return

    if key == "s":
        upper_actor = upper_actor.smoothLaplacian()
        lower_actor = lower_actor.smoothLaplacian()
        vp.interactor.Render()
        return

    if key =='d':
        upper_actor = upper_actor.decimate(1000 / upper_actor.N())
        lower_actor = lower_actor.decimate(1000 / lower_actor.N())
        vp.interactor.Render()
        return

    if key =='b':
        if u_actor:
            vp.remove(u_actor)

        u_actor = Actor(compute_u_plate(upper_actor.polydata(), lower_actor.polydata(), 20, margin=margin, plane_offset=plane_offset, resolution=args.res))
        u_actor.color('gray').alpha(0.5)
        vp.add(u_actor)
        if args.out:
            save(u_actor, args.out)
        vp.interactor.Render()
        return

def onMarginSlider(widget, event):
    global margin
    sliderRep = widget.GetRepresentation()
    margin = sliderRep.GetValue()

def onPlaneOffsetSlider(widget, event):
    global plane_offset
    sliderRep = widget.GetRepresentation()
    plane_offset = sliderRep.GetValue()

if __name__ == '__main__':
    scriptdir = os.path.dirname(sys.argv[0])
    default_upper = os.path.join(scriptdir, "/home/jack/work/data/medical/U board test/hsl III class middle/up.stl")
    default_lower = os.path.join(scriptdir, "/home/jack/work/data/medical/U board test/hsl III class middle/down.stl")
    default_final = os.path.join(scriptdir, "/home/jack/work/data/medical/U board test/hsl III class middle/final_1.stl")
    parser = argparse.ArgumentParser()
    parser.add_argument('upper', type=str, default=default_upper,
                        help='path of the upper teeth')
    parser.add_argument('lower', type=str, default=default_lower,
                        help='path of the lower teeth')
    parser.add_argument('out', type=str, default=default_final,
                        help='path of the output u plate')
    parser.add_argument('--res', type=float, default=0.1,
                        help='resolution of boolean operation')
    parser.add_argument('--margin', type=float, default=0,
                        help='margin of the u plate out of the teeth')
    parser.add_argument('--offset', type=float, default=0,
                        help='offset of the u plate middel plane')

    args = parser.parse_args()
    vp = Plotter()
    upper_actor = load(args.upper, c='red').clean()
    lower_actor = load(args.lower, c='blue').clean()
    u_actor = None
    margin = args.margin
    plane_offset = args.offset
    # final_actor = load(args.final, c='green')
    vp.keyPressFunction = onKeyPress
    vp.addSlider2D(onMarginSlider, -2.0, 3.0, value=margin, pos=3, title='margin')
    vp.addSlider2D(onPlaneOffsetSlider, -5.0, 3.0, value=plane_offset, pos=4, title='plane offset')
    vp.show([upper_actor, lower_actor])